# -*- coding: utf-8 -*-
import csv
import time
from multiprocessing import Process, Manager
from makeLabel import makeLabel
from industryClassLib import allIndustryClass3

# 配置
labelList = allIndustryClass3  # 打标数据集
num_processes = 20 # 并行数

def process_chunk(processIdx, chunk, results, startIndex, errorCompanyList):
    """
    处理一个块（chunk）的数据
    :param chunk: 一块公司名字列表
    :param results: 结果队列
    :param startIndex: 块起始索引
    :param progress_dict: 进度字典
    """
    global labelList
    agentMakeLabel = makeLabel()
    for relative_index, company_name in enumerate(chunk):
        # 调用计算流程
        absolute_index = startIndex + relative_index
        agentMakeLabel.website = ""
        try:
            certainty, r, explain, logs = agentMakeLabel.progressByBackCompanyName(labelList, company_name, QAFlag=True)
            result_str = str(r).replace('"', "'")
            results.put((absolute_index, company_name, agentMakeLabel.website, result_str, certainty))
        except Exception as e:
            errorCompanyList.append(company_name)
            print(f"【ERROR】程序运行错误 {company_name}: {str(e)}")
            results.put((absolute_index, company_name, agentMakeLabel.website, '', ''))
        # 更新进度信息
        print(f"进程{processIdx}({relative_index+1}/{len(chunk)})，{company_name}")
        print(logs)


def main():
    global labelList, num_processes
    t1 = time.time()
    # 读取文件
    with open("makeLabelBatchIn.txt", "r", encoding="utf-8") as f:
        lines = [line.strip() for line in f if line.strip()]

    # 初始化管理器
    manager = Manager()
    results = manager.Queue()
    errorCompanyList = manager.list()

    # 分配任务
    total_companies = len(lines)
    chunk_size = total_companies // num_processes + 1
    chunks = [lines[i:i + chunk_size] for i in range(0, total_companies, chunk_size)]
    print("总企业数：%s，拆分成%s进程并行计算，每个进程约%s个。"%(total_companies, num_processes, chunk_size))

    # 创建进程池并分配任务
    processes = []
    for i, chunk in enumerate(chunks):
        if len(chunk) == 0:
            continue
        start_index = i * chunk_size
        p = Process(target=process_chunk, args=(i, chunk, results, start_index, errorCompanyList))
        processes.append(p)
        p.start()

    # 等待所有进程完成
    for p in processes:
        p.join()

    # 准备一个列表来存放结果，按照索引排序
    sorted_results = [None] * total_companies

    # 按照原始顺序写入文件
    with open("makeLabelBatchOut.csv", "w", newline='', encoding="utf-8") as outFile:
        writer = csv.writer(outFile)
        writer.writerow(["companyName", "website", "result", "certainty"])
        # 从队列中获取结果并排序
        while not results.empty():
            index, company_name, website, result_str, certainty = results.get()
            sorted_results[index] = (company_name, website, result_str, certainty)
        # 写入文件
        for item in sorted_results:
            if item is not None:
                writer.writerow(item)

    if len(errorCompanyList) > 0:
        print("【计算错误企业】：", errorCompanyList)
    print("总企业数：%s，总用时：%s秒" % (total_companies, round(time.time() - t1, 2)))
    print("END")


if __name__ == "__main__":
    main()